﻿#pragma once
// Inner Fire 游戏引擎库
// 模糊类
//
// Copyright (c) 2023 by 尔西大帝. All rights reserved.
//
// 用于模糊的工具类。根据《DirectX12游戏开发实战》中的工具类修改。
// 原作者：Frank Luna
//
// Author: 尔西大帝
// Email: 2670613871@qq.com
// Created: 2024-11-26

#include "blur_filter.h"

namespace ifire {

BlurFilter::BlurFilter(
    ID3D12Device* device, UINT width, UINT height, DXGI_FORMAT format)
    : device_(device), width_(width), height_(height), format_(format) {

  BuildResources();
}

ID3D12Resource* BlurFilter::Output() { return blur_map_0_.Get(); }

void BlurFilter::Restore(
    ID3D12GraphicsCommandList* cmd, ID3D12Resource* back_buffer) {
  //
  // “倒腾”的第（四）步：
  //

  // 4.1 将back_buffer设为复制对象
  auto copy = CD3DX12_RESOURCE_BARRIER::Transition(back_buffer,
      D3D12_RESOURCE_STATE_COPY_SOURCE, D3D12_RESOURCE_STATE_COPY_DEST);
  cmd->ResourceBarrier(1, &copy);

  // 4.2 将生成的纹理复制到back_buffer
  cmd->CopyResource(back_buffer, Output());

  // 4.3 恢复到Render_Target的状态（如果有后续操作）
  auto restore = CD3DX12_RESOURCE_BARRIER::Transition(back_buffer,
      D3D12_RESOURCE_STATE_COPY_DEST, D3D12_RESOURCE_STATE_RENDER_TARGET);
  cmd->ResourceBarrier(1, &restore);
}

void BlurFilter::BuildDescriptors(CD3DX12_CPU_DESCRIPTOR_HANDLE hCpuDescriptor,
    CD3DX12_GPU_DESCRIPTOR_HANDLE hGpuDescriptor, UINT descriptorSize) {
  blur_0_cpu_srv_ = hCpuDescriptor;
  blur_0_cpu_uav_ = hCpuDescriptor.Offset(1, descriptorSize);
  blur_1_cpu_srv_ = hCpuDescriptor.Offset(1, descriptorSize);
  blur_1_cpu_uav_ = hCpuDescriptor.Offset(1, descriptorSize);

  blur_0_gpu_srv_ = hGpuDescriptor;
  blur_0_gpu_uav_ = hGpuDescriptor.Offset(1, descriptorSize);
  blur_1_gpu_srv_ = hGpuDescriptor.Offset(1, descriptorSize);
  blur_1_gpu_uav_ = hGpuDescriptor.Offset(1, descriptorSize);

  BuildDescriptors();
}

void BlurFilter::OnResize(UINT newWidth, UINT newHeight) {
  if ((width_ != newWidth) || (height_ != newHeight)) {
    width_ = newWidth;
    height_ = newHeight;

    // 需要重新构造资源
    BuildResources();
    BuildDescriptors();
  }
}

void BlurFilter::Execute(ID3D12GraphicsCommandList* cmd,
    ID3D12RootSignature* root_sig, ID3D12PipelineState* horz_blur_pso,
    ID3D12PipelineState* vert_blur_pso, ID3D12Resource* input, int blur_count) {
  // 用于代码更清晰的阅读
#define TRANS_ CD3DX12_RESOURCE_BARRIER::Transition
#define READ_ D3D12_RESOURCE_STATE_GENERIC_READ
#define ACCESS_ D3D12_RESOURCE_STATE_UNORDERED_ACCESS
#define COMMON_ D3D12_RESOURCE_STATE_COMMON
  // 计算高斯模糊权重（这个是纯数学，照公式抄）
  auto weights = CalcGaussWeights(2.5f);
  int blurRadius = (int)weights.size() / 2;

  // 设置计算着色器的根签名
  cmd->SetComputeRootSignature(root_sig);

  // 设置通用的常量
  cmd->SetComputeRoot32BitConstants(0, 1, &blurRadius, 0);
  cmd->SetComputeRoot32BitConstants(0, (UINT)weights.size(), weights.data(), 1);

  //
  // “倒腾”的第（一）步：
  //
  // 1.1将输入的项目设置为拷贝源
  auto src_barrier = TRANS_(input, D3D12_RESOURCE_STATE_RENDER_TARGET,
      D3D12_RESOURCE_STATE_COPY_SOURCE);
  cmd->ResourceBarrier(1, &src_barrier);

  // 1.2将map0设置为拷贝目标
  auto dest_barrier =
      TRANS_(blur_map_0_.Get(), COMMON_, D3D12_RESOURCE_STATE_COPY_DEST);
  cmd->ResourceBarrier(1, &dest_barrier);

  // 1.3 进行拷贝，input=>map0。
  cmd->CopyResource(blur_map_0_.Get(), input);

  // 1.4 将已经得到input数据的map0转为“可读”，用做计算着色器的输入。
  auto to_read =
      TRANS_(blur_map_0_.Get(), D3D12_RESOURCE_STATE_COPY_DEST, READ_);
  cmd->ResourceBarrier(1, &to_read);
  // 1.5 将map1转为“随机访问”。用做计算着色器的输出。
  auto set_access = TRANS_(blur_map_1_.Get(), COMMON_, ACCESS_);
  cmd->ResourceBarrier(1, &set_access);

  //
  // 进行模糊运算
  //
  for (int i = 0; i < blur_count; ++i) {
    //
    // 水平（Horizontal）模糊PASS
    //

    // 设置水平模糊的PSO
    cmd->SetPipelineState(horz_blur_pso);
    // 设置map0做为输入的纹理(SRV)，map1做为输出的纹理（UAV)
    cmd->SetComputeRootDescriptorTable(1, blur_0_gpu_srv_);
    cmd->SetComputeRootDescriptorTable(2, blur_1_gpu_uav_);

    // 计算我们需要派遣多少组来覆盖一行像素？其中每组覆盖256个像素。
    // 这个256是在ComputeShader中定义的。
    UINT num_groups_x = (UINT)ceilf(width_ / 256.0f);
    // 执行运算
    cmd->Dispatch(num_groups_x, height_, 1);

    //
    // “倒腾”的第（二）步：
    //
    // 2.1 将map0做为UAV，现在轮到他做为输出了。
    // 因为这时候他已经完成使命。计算结果在map1中。
    auto switch_access = TRANS_(blur_map_0_.Get(), READ_, ACCESS_);
    cmd->ResourceBarrier(1, &switch_access);

    // 2.2 将map1做为SRV，现在map1中的数据用于输入，计算垂直。
    auto switch_read = TRANS_(blur_map_1_.Get(), ACCESS_, READ_);
    cmd->ResourceBarrier(1, &switch_read);

    //
    // 垂直（Vertical）模糊PASS
    //

    // 设置垂直模糊的PSO
    cmd->SetPipelineState(vert_blur_pso);
    // 设置map1做为输入的纹理(SRV)，map0做为输出的纹理（UAV)
    cmd->SetComputeRootDescriptorTable(1, blur_1_gpu_srv_);
    cmd->SetComputeRootDescriptorTable(2, blur_0_gpu_uav_);

    // 计算多少组来覆盖一列像素。
    UINT num_groups_y = (UINT)ceilf(height_ / 256.0f);
    // 执行运算，注意此时用的是Dispatch的Y。和上面不同
    cmd->Dispatch(width_, num_groups_y, 1);

    //
    // “倒腾”的第（三）步：
    //

    // 3.1 将map0恢复为普通状态（SRV）。这是为了下次一运算做准备。
    // 因为需要执行blur_count次运算。
    auto restore_read = TRANS_(blur_map_0_.Get(), ACCESS_, READ_);
    cmd->ResourceBarrier(1, &restore_read);
    // 3.2 将map1变为输出状态（UAV）。
    auto restore_access = TRANS_(blur_map_1_.Get(), READ_, ACCESS_);
    cmd->ResourceBarrier(1, &restore_access);
  }

  //
  // “倒腾”的第（四）步：
  // 将项目输出到后备缓冲区
  //
  Restore(cmd, input);
  
  // “倒腾”的第（五）步：
  // 将状态复元为默认（因为下次调用这个函数时，初始时默认这两个为COMMOM）
  auto read_to_commom = TRANS_(blur_map_0_.Get(), READ_, COMMON_);
  cmd->ResourceBarrier(1, &read_to_commom);
  auto access_to_commom = TRANS_(blur_map_1_.Get(), ACCESS_, COMMON_);
  cmd->ResourceBarrier(1, &access_to_commom);
}

std::vector<float> BlurFilter::CalcGaussWeights(float sigma) {
  float twoSigma2 = 2.0f * sigma * sigma;

  // 根据sigma估算模糊半径，因为sigma控制了钟形曲线的“宽度”。
  // 例如，对于sigma = 3，钟形曲线的宽度是
  int blurRadius = (int)ceil(2.0f * sigma);

  assert(blurRadius <= MAX_BLUR_RADIUS);

  std::vector<float> weights;
  weights.resize(2 * blurRadius + 1);

  float weightSum = 0.0f;

  for (int i = -blurRadius; i <= blurRadius; ++i) {
    float x = (float)i;

    // 这个expf(-x * x / twoSigma2)就是著名的高斯模糊公式了。
    weights[i + blurRadius] = expf(-x * x / twoSigma2);

    weightSum += weights[i + blurRadius];
  }

  // 除以总和，使得所有权重加起来等于1.0。
  for (int i = 0; i < weights.size(); ++i) {
    weights[i] /= weightSum;
  }

  return weights;
}

void BlurFilter::BuildDescriptors() {
  D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
  srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
  srvDesc.Format = format_;
  srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
  srvDesc.Texture2D.MostDetailedMip = 0;
  srvDesc.Texture2D.MipLevels = 1;

  D3D12_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};

  uavDesc.Format = format_;
  uavDesc.ViewDimension = D3D12_UAV_DIMENSION_TEXTURE2D;
  uavDesc.Texture2D.MipSlice = 0;

  // “倒腾”的准备工作：
  // 每个map都需要两个状态，一个SRV，一个UAV。
  device_->CreateShaderResourceView(
      blur_map_0_.Get(), &srvDesc, blur_0_cpu_srv_);
  device_->CreateUnorderedAccessView(
      blur_map_0_.Get(), nullptr, &uavDesc, blur_0_cpu_uav_);

  device_->CreateShaderResourceView(
      blur_map_1_.Get(), &srvDesc, blur_1_cpu_srv_);
  device_->CreateUnorderedAccessView(
      blur_map_1_.Get(), nullptr, &uavDesc, blur_1_cpu_uav_);
}

void BlurFilter::BuildResources() {
  // 注意，压缩格式不能用于无序访问视图（UAV）。我们会得到类似这样的错误：
  // 错误：ID3D11Device::CreateTexture2D：格式（0x4d, BC3_UNORM）
  // 不能绑定为无序访问视图，也不能转换为可以绑定为无序访问视图的格式。
  // 因此，这种格式不支持D3D11_BIND_UNORDERED_ACCESS。

  D3D12_RESOURCE_DESC texDesc;
  ZeroMemory(&texDesc, sizeof(D3D12_RESOURCE_DESC));
  texDesc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
  texDesc.Alignment = 0;
  texDesc.Width = width_;
  texDesc.Height = height_;
  texDesc.DepthOrArraySize = 1;
  texDesc.MipLevels = 1;
  texDesc.Format = format_;
  texDesc.SampleDesc.Count = 1;
  texDesc.SampleDesc.Quality = 0;
  texDesc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN;
  texDesc.Flags = D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS;

  CD3DX12_HEAP_PROPERTIES default_heap(D3D12_HEAP_TYPE_DEFAULT);

  ThrowIfFailed(device_->CreateCommittedResource(&default_heap,
      D3D12_HEAP_FLAG_NONE, &texDesc, D3D12_RESOURCE_STATE_COMMON, nullptr,
      IID_PPV_ARGS(&blur_map_0_)));

  ThrowIfFailed(device_->CreateCommittedResource(&default_heap,
      D3D12_HEAP_FLAG_NONE, &texDesc, D3D12_RESOURCE_STATE_COMMON, nullptr,
      IID_PPV_ARGS(&blur_map_1_)));
}

} // namespace ifire
